import socket
import requests
import urllib3
import ping3


urllib3.disable_warnings()


class HeartBeat:
    def __init__(self, host, timeout=4, **extra):
        self.host = host
        self.timeout = timeout

    def heartbeat(self):
        pass
    

class ICMP(HeartBeat):
    """ping"""

    def heartbeat(self):
        try:
            resp = ping3.ping(self.host, timeout=self.timeout)
            if resp:
                return {'code': 0, 'error': ''}
            if resp is None:
                return {'code': -1, 'error': '探测超时'}
            return {'code': -1, 'error': '探测失败'}
        except Exception as e:
            return {'code': -1, 'error': '探测失败'}


class TCP(HeartBeat):
    """端口探测"""
    def __init__(self, host, port, timeout=10, **extra):
        self.port = port
        super(TCP, self).__init__(host=host, timeout=timeout, **extra)

    def heartbeat(self):
        try:
            with socket.create_connection((self.host, self.port), self.timeout) as sock:
                sock.shutdown(socket.SHUT_RDWR)
            return {'code': 0, 'error': ''}
        except socket.timeout as e:
            return {'code': -1, 'error': '探测超时'}
        except Exception as e:
            return {'code': -1, 'error': '探测失败'}


class Http(HeartBeat):
    """http服务存活"""
    DefaultPort = {'http': 80, 'https': 443}

    def _build_url(self, scheme, host, port=None, path=''):
        url = f"{scheme}://{host}"
        if port is not None:
            url += f':{port}' if self.DefaultPort.get(scheme) != port else ''
        if path:
            url += f'/{path}' if not path.startswith('/') else path
        return url

    def __init__(self, scheme, host, port=None, path='', desire_code=None, method='GET', timeout=90, **kwargs):
        super(Http, self).__init__(host=host, timeout=timeout)

        self.method = method
        self.url = self._build_url(scheme, host, port, path)
        self.desire_code = desire_code
        self.kwargs = kwargs

    def heartbeat(self):
        try:
            kwargs = {'timeout': self.timeout, 'verify': False, **self.kwargs}
            resp = requests.request(self.method, self.url, **kwargs)
            if not self.desire_code:  # 不校验状态码
                return {'code': 0, 'error': ''}
            if resp.status_code != self.desire_code:
                return {'code': -1, 'error': '状态码不匹配'}
            return {'code': 0, 'error': ''}
        except requests.exceptions.Timeout as e:
            return {'code': -1, 'error': '探测超时'}
        except requests.exceptions.RequestException as e:
            return {'code': -1, 'error': '探测失败'}


if __name__ == '__main__':
    print(ICMP('192.168.101.79').heartbeat())
    print(TCP('192.168.101.79', 80).heartbeat())
    print(Http('https', '192.168.101.79', '/', desire_code=200).heartbeat())
